from ultralytics import YOLO
import yaml
import os
import shutil
import torch

def train_model():
    # 加载预训练模型
    model = YOLO("yolov8s.pt")  # 使用小模型作为基础

    # 训练配置
    config = {
        "data": "data.yaml",
        "epochs": 100,
        "patience": 30,
        "imgsz": 640,
        "batch": 96 if torch.cuda.get_device_properties(0).total_memory > 24e9 else 64,  # 自动判断batch大小
        "name": "vehicle_detection_v2",
        "save": True,
        "save_period": 10,
        "device": "0",
        "lr0": 0.001,  # 初始学习率从0.01降到0.001
        "lrf": 0.01,   # 最终学习率从0.1降到0.01
        "optimizer": "AdamW",  # 替换默认SGD
        "warmup_epochs": 5,  # 延长热身期
        "augment": True,  # 启用更强数据增强
        "mixup": 0.1,     # 添加mixup增强
        "copy_paste": 0.1,  # 添加copy-paste增强
        "fliplr": 0.5,    # 水平翻转概率保持
        "degrees": 10.0,  # 旋转角度范围增大
        "shear": 2.0,     # 剪切变换增强
        "hsv_h": 0.015,   # 色相增强
        "hsv_s": 0.7,     # 饱和度增强
        "hsv_v": 0.4,     # 明度增强
        "translate": 0.2,  # 平移增强
        "scale": 0.9,     # 缩放增强
        "erasing": 0.4,   # 随机擦除
        "mosaic": 1.0,    # 启用mosaic
        "close_mosaic": 10,  # 最后10epoch关闭mosaic
        "amp": True,      # 保持自动混合精度
        "weight_decay": 0.05,  # 权重衰减增强
        "single_cls": False,  # 确保多类别训练
        "overlap_mask": True,
        "val": True,      # 启用验证
        "plots": True     # 保存训练曲线
    }

    # 开始训练
    results = model.train(**config)

    # 自动将最佳模型复制为model.pt
    best_pt_path = "runs/detect/vehicle_detection_v22/weights/best.pt"
    if os.path.exists(best_pt_path):
        shutil.copy(best_pt_path, "model.pt")
        print(f"✅ 最佳模型已保存为: model.pt")
    else:
        print(f"❌ 错误：未找到训练输出文件 {best_pt_path}")


if __name__ == "__main__":
    train_model()